import numpy as np

import tensorflow as tf
from tensorflow import keras
from functools import reduce
from config import NUM_ANCHORS, NUM_CLASSES

"""
compose functions in python
Reference: https://mathieularose.com/function-composition-in-python/
"""


def compose(*functions):
    return reduce(lambda f, g: lambda x: f(g(x)), functions, lambda x: x)


"""
*args : 位置参数，一个元祖
**kwargs： 一个字典
"""


def conv(*args, **kwargs):
    conv_kwargs = {'kernel_regularizer': keras.regularizers.l2(5e-4)}
    conv_kwargs['padding'] = 'valid' if kwargs.get('strides') == (2, 2) else 'same'
    conv_kwargs.update(kwargs)
    return keras.layers.Conv2D(*args, **conv_kwargs)


def DBL(input, *args, **kwargs):
    """

    :param input:
    :param args:
           filters,
           kernel_size,
           strides=(1, 1),
           padding='valid',
           data_format=None,
           dilation_rate=(1, 1),
           activation=None,
           use_bias=True,
    :param kwargs:
    :return:
    """
    x = conv(*args, **kwargs)(input)
    x = keras.layers.BatchNormalization()(x)
    x = keras.layers.LeakyReLU(alpha=0.1)(x)
    return x


def resblock_body(x, num_filters, num_resnet_blocks):
    # ((top_pad, bottom_pad), (left_pad, right_pad))
    x = keras.layers.ZeroPadding2D(padding=((1, 0), (1, 0)), data_format='channels_last')(x)
    x = DBL(x, num_filters, (3, 3), strides=(2, 2))
    for i in range(num_resnet_blocks):
        y = DBL(x, num_filters // 2, (1, 1))
        y = DBL(y, num_filters, (3, 3))
        x = keras.layers.Add()(
            [x, y])  # 写成这样x = keras.layers.Add()[x, y] 就会报错TypeError: 'Add' object is not subscriptable
    return x


def darknet_body(input):
    x = DBL(input, 32, (3, 3))
    x = resblock_body(x, 64, 1)
    x = resblock_body(x, 128, 2)
    x = resblock_body(x, 256, 8)
    x = resblock_body(x, 512, 8)
    x = resblock_body(x, 1025, 4)
    return x


def make_last_layers(input, num_filters, out_filters):
    x = DBL(input, num_filters, (1, 1))
    x = DBL(x, num_filters * 2, (3, 3))
    x = DBL(x, num_filters, (1, 1))
    x = DBL(x, num_filters * 2, (3, 3))
    x = DBL(x, num_filters, (1, 1))

    y = DBL(x, num_filters * 2, (3, 3))
    y = conv(out_filters, (1, 1))(y)

    return x, y


def yolo_body(inputs, num_anchors, num_classes):
    # inputs must be a tf.keras.Input对象
    darknet = keras.Model(inputs=inputs, outputs=darknet_body(inputs))
    x, y1 = make_last_layers(darknet.output, 512, num_anchors * (5 + num_classes))
    x = DBL(x, 256, (1, 1))
    x = keras.layers.UpSampling2D(2)(x)
    x = keras.layers.Concatenate()([x, darknet.layers[152].output])
    x, y2 = make_last_layers(x, 256, num_anchors * (5 + num_classes))
    x = DBL(x, 128, (1, 1))
    x = keras.layers.UpSampling2D(2)(x)
    x = keras.layers.Concatenate()([x, darknet.layers[92].output])
    x, y3 = make_last_layers(x, 128, num_anchors * (5 + num_classes))
    
    return keras.Model(inputs, [y1, y2, y3])


if __name__ == '__main__':
    x = np.random.random(size=(1, 416, 416, 3))
    # 判断一个变量是否为numpy数据类型
    if isinstance(x, np.ndarray):
        print(x.shape)
        # float64
        print(x.dtype)
        # 需要转化为float32，否则会报错为TypeError: Value passed to parameter 'x' has DataType float64 not in list of allowed values: float32
    # 转化为float32
    x = x.astype(np.float32)
    print(x.dtype)          # float32


    x = tf.keras.Input(shape=(416, 416, 3), dtype=tf.float32)
    # 这里要注意，传给kera.Model的inputs必须是一个tf的tensor数据类型
    model = yolo_body(x, NUM_ANCHORS, NUM_CLASSES)
    print(model.summary())
    print(model.input)
    for l in model.output:
        print(l.shape)
    # print(model.output)
    # print(model.layers[0])
    # print(model.layers[0].output)        # 用于网络构建的时候需要传每一层的".output"结果